from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial.distance import pdist, squareform
from scipy.stats import pearsonr

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal


def preprocess(trjs, max_length):
    processed_trjs = []
    for i in range(trjs.shape[0]):
        trj = trjs[i]
        if trj.shape[0] < max_length:
            last_element = trj[-1]
            processed_trj = np.concatenate([trj, np.repeat(last_element[np.newaxis,:], max_length - trj.shape[0], axis=0)], axis=0)
            processed_trjs.append(processed_trj)
        else:
            processed_trjs.append(trj)
    processed_trjs = np.array(processed_trjs)
    return processed_trjs
    
@jax.jit
def build_radiance_field(p):
    img = jnp.zeros((21, 21))
    top = 10
    bottom = 0
    effective_radius = 21*1.414
    px, py = p[0], p[1]

    x_coord_map = jnp.arange(img.shape[0])
    x_coord_map = jnp.repeat(x_coord_map[:, jnp.newaxis], img.shape[1], axis=1)

    y_coord_map = jnp.arange(img.shape[1])
    y_coord_map = jnp.repeat(y_coord_map[jnp.newaxis, :], img.shape[0], axis=0)
    
    dist = jnp.sqrt((x_coord_map - px)**2 + (y_coord_map - py)**2)
    img = jnp.where(dist < effective_radius,
                    (top - bottom) * (effective_radius - dist) / effective_radius + bottom,
                    img)
    return img

@jax.jit
def get_max_radiance_field(imgs):
    img = jnp.max(imgs, axis=0)
    return img

@jax.jit
def build_ridge(A):

    # 将 A 的第一个点对齐到 (10,10)
    A = A-A[0]+jnp.array([10,10])

    # 使用 jnp.map 来并行计算每个辐射场图像
    imgs = jax.vmap(build_radiance_field)(A)
    img = get_max_radiance_field(imgs)
    return img

build_ridge_vmap = jax.vmap(build_ridge)


def main():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration

    nn_type = ''
    if rpl_config.nn_type == "vanilla":
        nn_type = "vanilla"
    elif rpl_config.nn_type == "gru":
        nn_type = "gru"

    def load_data_and_compute(nn_type, seq_len, redundancy, diverse_set_capacity):

        rnn_limit_rings_file_name = "./logs/rnn_limit_rings_of_best_estimation_" + nn_type + "_" + str(seq_len) + "_" + str(redundancy) + "_" + str(diverse_set_capacity) + ".npz"
        # 载入 npz 文件
        rnn_limit_rings_of_best_estimation_file = np.load(rnn_limit_rings_file_name)

        # 获取 npz 文件中的所有对象名称
        matrix_names = rnn_limit_rings_of_best_estimation_file.files

        rnn_limit_rings_of_best_estimation = []

        # 遍历对象名称，访问和操作每个矩阵对象
        for name in matrix_names:
            matrix = rnn_limit_rings_of_best_estimation_file[name]
            # 在这里进行对矩阵对象的操作
            # 例如，打印矩阵的形状
            # print(f"Matrix '{name}' shape: {matrix.shape}")
            rnn_limit_rings_of_best_estimation.append(matrix)

        # 求 rnn_limit_rings_of_best_estimation 的中心位置序列
        rnn_limit_rings_of_best_estimation_center = []
        for i in range(len(rnn_limit_rings_of_best_estimation)):
            rnn_limit_rings_of_best_estimation_center.append(np.mean(rnn_limit_rings_of_best_estimation[i], axis=(0,1)))
        rnn_limit_rings_of_best_estimation_center = np.array(rnn_limit_rings_of_best_estimation_center)
        print("rnn_limit_rings_of_best_estimation_center.shape: ", rnn_limit_rings_of_best_estimation_center.shape)

        file_name = "obs_data_" + nn_type + "_" + str(seq_len) + "_" + str(redundancy) + "_" + str(diverse_set_capacity) + ".npz"
        obs_file = np.load("./logs/" + file_name)
        diverse_set_trajectoies = obs_file["diverse_set_trajectoies"]
        print("diverse_set_trajectoies.shape: ", diverse_set_trajectoies.shape)

        return rnn_limit_rings_of_best_estimation_center, diverse_set_trajectoies
    
    configs = [

        [nn_type, 6, 1, 100],
        [nn_type, 7, 1, 100],
        [nn_type, 8, 1, 100],
        [nn_type, 9, 1, 100],
        [nn_type, 10, 1, 100],
        [nn_type, 11, 1, 100],
        [nn_type, 12, 1, 100],
        [nn_type, 13, 1, 100],
        [nn_type, 14, 1, 100],
        [nn_type, 15, 1, 100],
        
        ]

    rnn_limit_rings_of_best_estimation_centers = []
    ring_lengths = []
    diverse_set_trajectoies = []
    for i in range(len(configs)):
        _rnn_limit_rings_of_best_estimation_centers, dts = load_data_and_compute(configs[i][0], configs[i][1], configs[i][2], configs[i][3])
        rnn_limit_rings_of_best_estimation_centers.append(_rnn_limit_rings_of_best_estimation_centers)
        diverse_set_trajectoies.append(dts)
        ring_length = configs[i][1]
        data_cap = _rnn_limit_rings_of_best_estimation_centers.shape[0]
        ring_length_ = np.array([ring_length for i in range(data_cap)])
        ring_lengths.append(ring_length_)

    ring_lengths = np.array(ring_lengths)
    ring_lengths = ring_lengths.reshape(ring_lengths.shape[0]*ring_lengths.shape[1])
    print(ring_lengths)

    diverse_set_trajectoies_mat = []
    for i in range(len(diverse_set_trajectoies)):
        length_aligned_trj = preprocess(diverse_set_trajectoies[i], 15)
        diverse_set_trajectoies_mat.append(length_aligned_trj)
        # print("diverse_set_trajectoies_mat[i].shape: ", diverse_set_trajectoies_mat[i].shape)
    diverse_set_trajectoies_mat = np.concatenate(diverse_set_trajectoies_mat, axis=0)
    print("diverse_set_trajectoies_mat.shape: ", diverse_set_trajectoies_mat.shape)

    # 将 rnn_limit_rings_of_best_estimation_centers 所有元素拼接起来
    rnn_limit_rings_of_best_estimation_center_mat = np.concatenate(rnn_limit_rings_of_best_estimation_centers, axis=0)
    print("rnn_limit_rings_of_best_estimation_center_mat.shape: ", rnn_limit_rings_of_best_estimation_center_mat.shape)

    ridge_images = build_ridge_vmap(diverse_set_trajectoies_mat)
    print("RIs.shape: ", ridge_images.shape)
    # 将 RIs 从 (100, 21, 21) 转换成 (100, 441)
    ridge_images = ridge_images.reshape(ridge_images.shape[0], ridge_images.shape[1] * ridge_images.shape[2])
    print("RIs.shape: ", ridge_images.shape)
    # trj_dist_matrix = squareform(pdist(ridge_images, 'euclidean'))
    
    # 对 rnn_limit_rings_of_best_estimation_center_mat 进行 PCA
    pca = PCA()
    pca.fit(rnn_limit_rings_of_best_estimation_center_mat)
    rnn_limit_rings_of_best_estimation_center_mat_pca = pca.transform(rnn_limit_rings_of_best_estimation_center_mat)
    # print("rnn_limit_rings_of_best_estimation_center_mat_pca.shape: ", rnn_limit_rings_of_best_estimation_center_mat_pca.shape)
    # # rnn_limit_rings_of_best_estimation_center_mat_pca[:,16:] = 0
    # rnn_limit_rings_of_best_estimation_center_mat = pca.inverse_transform(rnn_limit_rings_of_best_estimation_center_mat_pca)

    rnn_limit_rings_of_best_estimation_center_mat_pca_3d = rnn_limit_rings_of_best_estimation_center_mat_pca[:, :2]
    # 计算 rnn_limit_rings_of_best_estimation_center_mat_pca_3d 中每个点相对于中心点的角度
    rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles = []
    for i in range(rnn_limit_rings_of_best_estimation_center_mat_pca_3d.shape[0]):
        rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.append(np.arctan2(rnn_limit_rings_of_best_estimation_center_mat_pca_3d[i,1], rnn_limit_rings_of_best_estimation_center_mat_pca_3d[i,0]))
    rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles = np.array(rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles)
    # 将 rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles 转换到 [0, 2pi] 区间
    rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles = rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles - rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.min()
    # 打印出 rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles 的最大值和最小值
    print("rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.max(): ", rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.max())
    print("rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.min(): ", rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.min())

    """ random sample
    """

    corr_lpf = 0

    while True:
    # for ii in range(1):

        # 在 最大值和最小值之间随机选中一个角度
        random_angle = random.uniform(rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.min(), rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.max())
        sector_angle = np.pi / 32
        # print("random_angle: ", random_angle)
        # 选中所有角度在 [random_angle - sector_angle, random_angle + sector_angle] 之间的点
        idx = np.where((rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles >= random_angle - sector_angle) & (rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles <= random_angle + sector_angle))
        # print("idx", idx)
        rnn_limit_rings_of_best_estimation_center_mat_selected = rnn_limit_rings_of_best_estimation_center_mat[idx]
        ridge_images_selected = ridge_images[idx]
        seq_len_selected = ring_lengths[idx]
        rnn_limit_rings_of_best_estimation_center_mat_selected_pca = rnn_limit_rings_of_best_estimation_center_mat_pca[idx]

        # 对 rnn_limit_rings_of_best_estimation_center_mat_pca 进行 3D 可视化
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.scatter(
            rnn_limit_rings_of_best_estimation_center_mat_pca[:,0], 
            rnn_limit_rings_of_best_estimation_center_mat_pca[:,1], 
            rnn_limit_rings_of_best_estimation_center_mat_pca[:,2], 
            c = 'b',
            s=5)
        ax.scatter(
            rnn_limit_rings_of_best_estimation_center_mat_selected_pca[:,0], 
            rnn_limit_rings_of_best_estimation_center_mat_selected_pca[:,1], 
            rnn_limit_rings_of_best_estimation_center_mat_selected_pca[:,2], 
            c = 'r',
            s=50)
        ax.view_init(elev=90, azim=0)
        plt.show()

        # # 获取 seq_len_selected 按照长度从小到大排序的索引
        # idx_reodered = np.argsort(seq_len_selected)
        # rnn_limit_rings_of_best_estimation_center_mat_selected = rnn_limit_rings_of_best_estimation_center_mat_selected[idx_reodered]
        # ridge_images_selected = ridge_images_selected[idx_reodered]
        # seq_len_selected = seq_len_selected[idx_reodered]
        # print(seq_len_selected)

        # """ 从 ridge_images 中采样和 ridge_images_selected 相同数量的样本, 用作随机对照组
        # """
        # idx_rnd = np.random.choice(ridge_images.shape[0], ridge_images_selected.shape[0], replace=False)
        # ridge_images_selected = ridge_images[idx_rnd]

        trj_dist_matrix = squareform(pdist(ridge_images_selected, 'euclidean'))
        limit_ring_dist_matrix = squareform(pdist(rnn_limit_rings_of_best_estimation_center_mat_selected, 'euclidean'))

        def del_diag(matrix):
            new_matrix = np.zeros((matrix.shape[0], matrix.shape[1] - 1))
            for i in range(matrix.shape[0]):
                for j in range(matrix.shape[1]):
                    if i != j:
                        new_matrix[i, j if j < i else j - 1] = matrix[i, j]
            return new_matrix
        
        trj_dist_matrix = del_diag(trj_dist_matrix)
        limit_ring_dist_matrix = del_diag(limit_ring_dist_matrix)

        trj_dist_matrix_norm = trj_dist_matrix / np.max(trj_dist_matrix)
        limit_ring_dist_matrix_norm = limit_ring_dist_matrix / np.max(limit_ring_dist_matrix)

        """ compute sim matrix
        """
        # 将两个自相似矩阵转换成两个向量，然后计算两个向量的pearson相关系数
        trj_dist_matrix_vector = trj_dist_matrix_norm.reshape(-1)
        limit_ring_dist_matrix_vector = limit_ring_dist_matrix_norm.reshape(-1)
        # 计算 pearson 相关系数
        corr, _ = pearsonr(trj_dist_matrix_vector, limit_ring_dist_matrix_vector)
        # corr = np.corrcoef(trj_dist_matrix_norm, limit_ring_dist_matrix_norm)
        # print("corr.shape: ", corr.shape)
        # corr = np.sqrt(np.sum(corr**2))

        corr_lpf = 0.95*corr_lpf + 0.05*corr
        print("corr_lpf: ", corr_lpf)

        """ 显示两个自相似矩阵
        """
        # 将两个自相似矩阵合并在同一个视图中,使用两个subplot进行可视化
        fig = plt.figure(figsize=(10, 10))
        ax1 = fig.add_subplot(121)
        ax1.imshow(trj_dist_matrix_norm)
        for i in range(seq_len_selected.shape[0]):
            ax1.text(-1, i, str(seq_len_selected[i]), ha='center', va='center')

        ax2 = fig.add_subplot(122)
        ax2.imshow(limit_ring_dist_matrix_norm)
        # 添加子图标题以及颜色条
        ax1.set_title(" trj_dist_matrix")
        ax2.set_title(" limit_ring_dist_matrix")
        fig.colorbar(ax1.imshow(trj_dist_matrix_norm), ax=ax1)
        fig.colorbar(ax2.imshow(limit_ring_dist_matrix_norm), ax=ax2)

        # 在第一个图表下方添加文本
        ax1.text(-0.2, 0.3, f"pearson_correlation: {corr:.5f}", color='red', ha='center', va='center',fontweight='bold', fontsize=12, transform=ax1.transAxes)

        plt.show()



if __name__ == "__main__":
    main()